import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import shapefile as shp
import seaborn as sns


def read_shapefile(sf):
    #读取shp文件
    fields = [x[0] for x in sf.fields][1:]
    records = [list(i) for i in sf.records()]
    shps = [s.points for s in sf.shapes()]
    df = pd.DataFrame(columns=fields, data=records)
    df = df.assign(coords=shps)
    return df


def plot_shape(id, s=None):
    # 绘制地理行政区划的形状
    plt.figure()
    ax = plt.axes()
    ax.set_aspect('equal')
    shape_ex = sf.shape(id)
    x_lon = np.zeros((len(shape_ex.points), 1))
    y_lat = np.zeros((len(shape_ex.points), 1))
    for ip in range(len(shape_ex.points)):
        x_lon[ip] = shape_ex.points[ip][0]
        y_lat[ip] = shape_ex.points[ip][1]
    plt.plot(x_lon, y_lat)
    x0 = np.mean(x_lon)
    y0 = np.mean(y_lat)
    plt.text(x0, y0, s, fontsize=10)
    plt.xlim(shape_ex.bbox[0], shape_ex.bbox[2])
    return x0, y0


def plot_map(sf, x_lim=None, y_lim=None, figsize=(11, 9)):
    # 绘制特定形状的区域
    plt.figure(figsize=figsize)
    id = 0
    for shape in sf.shapeRecords():
        x = [i[0] for i in shape.shape.points[:]]
        y = [i[1] for i in shape.shape.points[:]]
        plt.plot(x, y, 'k')

        if (x_lim == None) & (y_lim == None):
            x0 = np.mean(x)
            y0 = np.mean(y)
            plt.text(x0, y0, id, fontsize=10)
        id = id + 1

    if (x_lim != None) & (y_lim != None):
        plt.xlim(x_lim)
        plt.ylim(y_lim)


def plot_map_fill(id, sf, x_lim=None, y_lim=None, figsize=(11, 9), color='r'):
    # 对特定的形状进行颜色填充
    plt.figure(figsize=figsize)
    fig, ax = plt.subplots(figsize=figsize)
    for shape in sf.shapeRecords():
        x = [i[0] for i in shape.shape.points[:]]
        y = [i[1] for i in shape.shape.points[:]]
        ax.plot(x, y, 'k')
    shape_ex = sf.shape(id)
    x_lon = np.zeros((len(shape_ex.points), 1))
    y_lat = np.zeros((len(shape_ex.points), 1))
    for ip in range(len(shape_ex.points)):
        x_lon[ip] = shape_ex.points[ip][0]
        y_lat[ip] = shape_ex.points[ip][1]
    ax.fill(x_lon, y_lat, color)

    if (x_lim != None) & (y_lim != None):
        plt.xlim(x_lim)
        plt.ylim(y_lim)


def plot_map_fill_multiples_ids(title, city, sf, x_lim=None, y_lim=None, figsize=(11, 9), color='r'):
    #对特定的区块进行颜色填充
    plt.figure(figsize=figsize)
    fig, ax = plt.subplots(figsize=figsize)
    fig.suptitle(title, fontsize=16)
    for shape in sf.shapeRecords():
        x = [i[0] for i in shape.shape.points[:]]
        y = [i[1] for i in shape.shape.points[:]]
        ax.plot(x, y, 'k')

    for id in city:
        shape_ex = sf.shape(id)
        x_lon = np.zeros((len(shape_ex.points), 1))
        y_lat = np.zeros((len(shape_ex.points), 1))
        for ip in range(len(shape_ex.points)):
            x_lon[ip] = shape_ex.points[ip][0]
            y_lat[ip] = shape_ex.points[ip][1]
        ax.fill(x_lon, y_lat, color)

        x0 = np.mean(x_lon)
        y0 = np.mean(y_lat)
        plt.text(x0, y0, id, fontsize=10)

    if (x_lim != None) & (y_lim != None):
        plt.xlim(x_lim)
        plt.ylim(y_lim)


def plot_cities_2(sf, title, cities, color):
    df = read_shapefile(sf)
    city_id = []
    for i in cities:
        city_id.append(df[df.DIST_NAME == i.upper()].index[0])
        plot_map_fill_multiples_ids(title, city_id, sf, x_lim=None,
                                    y_lim=None, figsize=(11, 9), color=color)


def calc_color(data, color=None):
    # 颜色设置选择
    if color == 1:
        color_sq = ['#dadaebFF', '#bcbddcF0', '#9e9ac8F0', '#807dbaF0', '#6a51a3F0', '#54278fF0']
        colors = 'Purples'
    elif color == 2:
        color_sq = ['#c7e9b4', '#7fcdbb', '#41b6c4', '#1d91c0', '#225ea8', '#253494']
        colors = 'YlGnBu'
    elif color == 3:
        color_sq = ['#f7f7f7', '#d9d9d9', '#bdbdbd', '#969696', '#636363', '#252525']
        colors = 'Greys'
    elif color == 9:
        color_sq = ['#ff0000', '#ff0000', '#ff0000', '#ff0000', '#ff0000', '#ff0000']
    else:
        color_sq = ['#ffffd4', '#fee391', '#fec44f', '#fe9929', '#d95f0e', '#993404']
        colors = 'YlOrBr'
    new_data, bins = pd.qcut(data, 6, retbins=True, labels=list(range(6)))
    color_ton = []
    for val in new_data:
        color_ton.append(color_sq[val])
    if color != 9:
        colors = sns.color_palette(colors, n_colors=6)
        sns.palplot(colors, 0.6)
        for i in range(6):
            print("\n" + str(i + 1) + ': ' + str(int(bins[i])) +
                  " => " + str(int(bins[i + 1]) - 1))
        print("\n\n   1   2   3   4   5   6")
    return color_ton, bins


def plot_cities_data(sf, title, cities, data=None, color=None, print_id=False):
    # 绘制区域
    color_ton, bins = calc_color(data, color)
    df = read_shapefile(sf)
    city_id = []
    for i in cities:
        city_id.append(df[df.DIST_NAME == i.upper()].index[0])
    plot_map_fill_multiples_ids_tone(sf, title, city_id, print_id, color_ton,
                                     bins, x_lim=None, y_lim=None, figsize=(11, 9))


def plot_map_fill_multiples_ids_tone(sf, title, city, print_id, color_ton,
                                     bins, x_lim=None, y_lim=None, figsize=(11, 9)):
    plt.figure(figsize=figsize)
    fig, ax = plt.subplots(figsize=figsize)
    fig.suptitle(title, fontsize=16)
    for shape in sf.shapeRecords():
        x = [i[0] for i in shape.shape.points[:]]
        y = [i[1] for i in shape.shape.points[:]]
        ax.plot(x, y, 'k')

    for id in city:
        shape_ex = sf.shape(id)
        x_lon = np.zeros((len(shape_ex.points), 1))
        y_lat = np.zeros((len(shape_ex.points), 1))
        for ip in range(len(shape_ex.points)):
            x_lon[ip] = shape_ex.points[ip][0]
            y_lat[ip] = shape_ex.points[ip][1]
        ax.fill(x_lon, y_lat, color_ton[city.index(id)])                 # 这里将对应的颜色类型进行返回
        if print_id != False:
            x0 = np.mean(x_lon)
            y0 = np.mean(y_lat)
            plt.text(x0, y0, id, fontsize=10)
    if (x_lim != None) & (y_lim != None):
        plt.xlim(x_lim)
        plt.ylim(y_lim)


sns.set(style='whitegrid', palette='pastel', color_codes=True)
sns.mpl.rc('figure', figsize=(10, 6))
shp_path = r'D:\Program\轨迹数据挖掘\TowardsDataScience\Data\District_boundary\District_Boundary.shp'
# shp_path = r'1.shp'
sf = shp.Reader(shp_path)
df = read_shapefile(sf)
census_17 = df.POPULATION
census_17.shape
#plotting
title = 'Population Distrubution on Rajasthan Region'
data = census_17
names = df.DIST_NAME
plot_cities_data(sf, title, names, data, 1, True)

plt.savefig('1.png')
